//
// CollinsFragmentation.cxx
//
// Created by TB on 11/30/11.
// Copyright 2011 BNL. All rights reserved.
//

#include <cstdlib> // For abs(int)
#include <iostream>
#include <sstream>

#include <boost/foreach.hpp>

#include <TCanvas.h>
#include <TH1D.h>
#include <TH2D.h>
#include <TMath.h>
#include <TRandom.h>
#include <TRotation.h>

#include "tppmc/CollinsFragmentation.h"
#include "tppmc/Enumerations.h"
#include "tppmc/KkpParton.h"
#include "tppmc/TF2J.h"

extern "C" {
   
   // Access to FORTRAN fDSS subroutine.
   void fdss_(int* hadron, int* charge, int* order, double* z,
              double* scale, int* parton, double* result);
   
   // Access to FORTRAN FRAGINI COMMON block.
   extern struct {
      double fini; 
   } fragini_;
   
} // extern "C"

namespace {

double collins_z_dependence(double z, double a, double b, double n){
   return
   n *
   pow(z, a) *
   pow((1. - z), b) *
   pow((a + b), (a + b)) /
   (pow(a, a) * pow(b, b));
}

// Look for better alternative than global pointer
PARAMETERS* gParams;

/**
 2D function for DSS + Collins fragmentation.
 Wrapper around Fortran DSS subroutine (see fDSS.f) compatible with ROOT TF1.
 x[0] is the z value at which to evaluate the function.
 x[1] is the phi value.
 Depends on eight parameters, par[0] to par[7]:
 0: Average quark pT^2
 1: Q at which to evaluate
 2: Parton type, following KKP convention (see fDSS.f)
 3: LO or NLO function
 4: Hadron species (pion, kaon, proton or charged hadrons)
 5: Hadron charge (+, - or average of +/-)
 6: Depolarisation factor
 7: Transverse parton spin
 8: Collins parameterisation up quark M^2
 9: Collins parameterisation down quark M^2
 */
double DSSFragFunctionCollins(Double_t* x, Double_t* par) {//const {
                                                           // This FF is only valid for PIPLUS, PIMINUS or PIZERO
                                                           // It should really only be used for u,ubar,d,dbar,s,sbar quarks
                                                           // (otherwise use DSSFragFunction).
   int ih = static_cast<int>(par[4]);
   int iset = 0;            // parametrization is LO ONLY
   int charge = static_cast<int>(par[5]);    // parametrization is PIPLUS, PIMINUS,
                                             //PIZERO selected below...
   double z = x[0];
   double phi = x[1];
   double qs2 = par[1]*par[1];
   double dh=0;
   double dh2=0; //include disfavored FF for p0 as (fav + disfav)
   int qtyp = static_cast<int>(par[2]);
   double depol_factor = par[6];
   bool atPosivitityBound = static_cast<bool>(par[7]);
   double collins = 0.0; 
   static bool first = true;
   if(first) {
      fragini_.fini = 0;
      first = false; 
   } // if
     // Check to make sure we weren't called in some strange way. 
   if(par[4] not_eq 1.) {
      cerr <<
      "WARNING - DSSFragFunctionCollins called with par[4] = " <<
      par[4] << " (not PION, PION assumed)" << endl;
      ih = 1;
   } // if
   
   if((qtyp not_eq UQUARK) and
      (qtyp not_eq UBAR) and
      (qtyp not_eq DQUARK) and
      (qtyp not_eq DBAR) and
      (qtyp not_eq SQUARK) and
      (qtyp not_eq SBAR)) {
      qtyp = SQUARK; 
   } // if
   
   // Favored/Unfavored swaps: 
   // These lines enforces the favored/unfavored 
   // fragmentation equalities (see hep-ph0701006v3 eq.45,46):
   
   // u -> pi+   : FAVORED
   // dbar-> pi+ : FAVORED
   // d -> pi-   : FAVORED
   // ubar-> pi- : FAVORED
   
   // u -> pi0   : FAVORED + DISFAVORED
   // dbar-> pi0 : FAVORED + DISFAVORED
   // d -> pi0   : FAVORED + DISFAVORED
   // ubar-> pi0 : FAVORED + DISFAVORED
   
   // d -> pi+    : UNFAVORED
   // u -> pi-    : UNFAVORED
   // ubar -> pi+ : UNFAVORED
   // dbar -> pi- : UNFAVORED
   
   // s ->pi+/-/0     : UNFAVORED
   // sbar -> pi+/-/0 : UNFAVORED
   
   
   // The fragmentation functions are parametrized in terms of the fragmentation
   // function for a positive pion, and we choose the favored or unfavored fragmentation 
   // function from this by setting the quark type. 
   
   //      enum { FAVOURED, UNFAVOURED };
   
   int fqtyp = qtyp; 
   
   if(charge==1){
      // positive pion
      if(qtyp==UQUARK) fqtyp = UQUARK; 
      if(qtyp==DQUARK) fqtyp = DQUARK;
      if(qtyp==UBAR) fqtyp = DQUARK;
      if(qtyp==DBAR) fqtyp = UQUARK; 
      if((qtyp==SBAR)||(qtyp==SQUARK)) fqtyp = DQUARK; // always unfavored
   }
   else if(charge==-1){
      // negative pion
      if(qtyp==DQUARK) fqtyp = UQUARK;
      if(qtyp==UQUARK) fqtyp = DQUARK; 
      if(qtyp==UBAR) fqtyp = UQUARK;
      if(qtyp==DBAR) fqtyp = DQUARK; 
      if((qtyp==SBAR)||(qtyp==SQUARK)) fqtyp = DQUARK; // always unfavored
   }
   else if(charge==0){
      // neutral pion
      if(qtyp==UQUARK) fqtyp = UQUARK; 
      if(qtyp==DQUARK) fqtyp = UQUARK;
      if(qtyp==UBAR) fqtyp = UQUARK;
      if(qtyp==DBAR) fqtyp = UQUARK; 
      
      if((qtyp==SBAR)||(qtyp==SQUARK))
         fqtyp = DQUARK; // always unfavored
      else { // add disfavored piece to pi0
         int fqtyp2 = DQUARK; 
         int icp = 1; 
         fdss_(&ih,&icp,&iset,&z,&qs2,&fqtyp2,&dh2);
         dh2/=2;
      }
      
   }
   
   // Always called for PIPLUS, either UQUARK (favored)
   // or DQUARK (unfavored)
   int icp = 1; 
   fdss_(&ih,&icp,&iset,&z,&qs2,&fqtyp,&dh);
   
   if(charge==0&&fqtyp == UQUARK){
      //make sure, that pi0 for light quarks is 1/2 favored and 1/2 disfavored
      dh/=2;
   }
   
   // Protection from negative values...
   
   if(dh < 0.0) dh = 0.0; 
   
   const VALUES& params = gParams->Collins.parameters;
   
   // Need to calculate and factor in the Collins part
   // of the fragmentation function
   double ptq2_average = gParams->GetPtq2Average();
   double collins_width, integrand; 
   
   switch(fqtyp) {
         
      case(UQUARK): // FAVORED:
         if(atPosivitityBound){
            collins = 2.*dh;
         }
         else {
            
            collins = 2.*dh*collins_z_dependence(z, params.a_up, 
                                                 params.b_up, 
                                                 params.n_up);
            //               collins_width = params.m2_up*ptq2_average/(ptq2_average + params.m2_up);
            // This is only z-dependent Collins, so integrate out pT
            collins_width = params.m2_up * ptq2_average/(ptq2_average + params.m2_up);
            integrand = sqrt( TMath::Pi() * collins_width )/2./2.; // Mulders' definition
            integrand *=  4.; // See Eq (45.) Delta D = 4 * H(1/2) 
            
            collins *= integrand; 
         } // else
         break;
         
      case(DQUARK): // UNFAVORED:
         if(atPosivitityBound){
            collins = -2.*dh; 
         }
         else {
            collins = 2.*dh*collins_z_dependence(z, params.a_down, 
                                                 params.b_down, 
                                                 params.n_down);
            //               collins_width = params.m2_down*ptq2_average/(ptq2_average + params.m2_down);
            // This is only z-dependent Collins, so integrate out pT
            collins_width = params.m2_down * ptq2_average / (ptq2_average + params.m2_down);
            integrand = sqrt( TMath::Pi() * collins_width )/2./2.; // Mulders' definition
            integrand *=  4.; // See Eq (45.) Delta D = 4 * H(1/2) 
            
            collins *= integrand; 
         }
         break;
         
      default:
         collins = 0.0; 
         break;
   }
   
   if(charge==0){
      if(atPosivitityBound){
         collins += -2.*dh2; 
      }
      else {
         double collins2=0;
         collins2 = 2.*dh2*collins_z_dependence(z, params.a_down, 
                                                params.b_down, 
                                                params.n_down);
         //               collins_width = params.m2_down*ptq2_average/(ptq2_average + params.m2_down);
         // This is only z-dependent Collins, so integrate out pT
         collins_width = params.m2_down * ptq2_average / (ptq2_average + params.m2_down);
         integrand = sqrt( TMath::Pi() * collins_width )/2./2.; // Mulders' definition
         integrand *=  4.; // See Eq (45.) Delta D = 4 * H(1/2) 
         
         collins += collins2 * integrand; 
      }
      
   } //neutral pions 
   
   // Why is there no "spin" factor (sign) in this equation? Because what we want is the angular 
   // distribution referenced to the spin direction. This is accounted for properly 
   // when we calculate the angular distrubution of the hadrons in the fragmentation
   // because it is done with respect to the spin direction.
   
   // NOTE ALSO: The angular distribution that is given is with respect to the vector
   // p^ x jT^, so the dot product gives a cos(phi) behavior. Because this is perp. to
   // jt^ the hadron distribution will have a sin(phi) modulation.
   double full_FF = dh + dh2 + (0.5*cos(phi)*depol_factor*collins); 
   return std::max(full_FF, 0.);
}

} // anonymous namespace

namespace tppmc {

double CollinsFragmentation::Evaluate(const pythpart& parton,
                                      const Hadron& hadron,
                                      double scale) {
   if(not Supports(parton)) {
      return mDss.Evaluate(parton, hadron, scale);
   } // if
   
   switch(hadron.Type()) {
      case  211:
      case -211: // Deliberate fall-through
      case  111: // Deliberate fall-through
         break;
      default:
         return mDss.Evaluate(parton, hadron, scale);
         break;
   } // switch(h)
   
   Function()->SetParameter(SCALE,  sqrt(scale));
   Function()->SetParameter(PARTON, kkp::fromPdg(parton.type));
   mFunction->SetParameter(6, parton.cos_spin_depol);
   
   SetFunctionParameters(hadron.type);
   return Function()->Eval(hadron.z, hadron.PxPyPz().Phi());
}

CollinsFragmentation::CollinsFragmentation(const FragmentationFunction* ff,
                                           TVector3 spin,
                                           int order)
: mUnpolarised(ff)
, mSpin(spin.Unit()) // Normalise the spin vector
, mAtPositivityBound(false)
, mDss(order, 0.2)
, mFunction(NULL) {
   mParameters.Collins.SetReadFile("tables/collins_parameters.dat");
   mParameters.Collins.ReadParameters();
   gParams = &mParameters;
}

CollinsFragmentation::~CollinsFragmentation() {
   if(mFunction) {
      delete mFunction;
      mFunction = NULL;
   } // if
}

bool CollinsFragmentation::Initialise() {
   // Call Initialise() for the stored DssFragmentation.
   bool initialised = mDss.Initialise();
   // This inherits from DssFragmentation so call the base
   // class' Initialise().
   return initialised and DssFragmentation::Initialise();
}

TF2* CollinsFragmentation::Function() {
   if(not mFunction) {
      mFunction = new TF2J("tppmc_collins_frag_function",
                           DSSFragFunctionCollins,
                           mZLower, mZUpper,
                           -TMath::Pi(), TMath::Pi(),
                           8);
      // Set default parameters:
      mFunction->SetParameter(0, 1); // This isn't actually used any more!
      mFunction->SetParameter(1, 2.5); // Scale Q (not Q^2)
      mFunction->SetParameter(2, UQUARK); // Parton type
      mFunction->SetParameter(3, LO); // Calculation order
      mFunction->SetParameter(4, 1); // Hadron type
      mFunction->SetParameter(5, 1); // Charge
      mFunction->SetParameter(6, 1.); // Depolarisation factor
      mFunction->SetParameter(7, mAtPositivityBound);
      // Set number of points in x, y so TF2 interpolation is better
      mFunction->SetNpx(500);
      mFunction->SetNpy(50);
      mFunction->SetRange(mZLower,-TMath::Pi(), mZUpper,TMath::Pi());
   } // if
   return mFunction;
}

void CollinsFragmentation::Qa(std::string filename) {
   TCanvas canvas;
   canvas.Print((filename + "[").c_str());
   int partons[] = { 2, 1 };
   int hadrons[] = { 211, -211, 111 };
   float depols[] = { 0.1, 0.3, 0.5, 0.7 };
   std::ostringstream oss;
   BOOST_FOREACH(int parton, partons) {
      BOOST_FOREACH(int hadron, hadrons) {
         BOOST_FOREACH(double depol, depols) {
            SetFunctionParameters(hadron);
            Function()->SetParameter(1, 10.);
            Function()->SetParameter(2, kkp::fromPdg(parton));
            Function()->SetParameter(6, depol);
            oss << "Parton " << parton << " hadron " << hadron << " depol " << depol;
            TH1D hist("", oss.str().c_str(), 9, 0.05, 0.95);
            oss.str("");
            for(int i(1); i <= 9; ++i) {
               double integral = Function()->Integral(0.05, 0.05 + i * 0.1);
               hist.SetBinContent(i, integral);
            } // for
            hist.Scale(1. / hist.GetBinContent(hist.GetNbinsX()));
            hist.SetMinimum(0.);
            hist.Draw();
            canvas.Print(filename.c_str());
         }
      }
   }
   canvas.Print((filename + "]").c_str());
}

std::vector<Hadron>
CollinsFragmentation::Fragment(const pythpart& parton, double scale) {
   
   if(not Supports(parton)) {
      return mDss.Fragment(parton, scale);
   } // if
     //note that the variable scale contains pt2 or Q2, but Parameter 1 needs to be Q instead
   double pt = std::max(sqrt(scale), sqrt(MinimumScale()));
   
   
   Function()->SetParameter(1, pt);
   Function()->SetParameter(2, kkp::fromPdg(parton.type));
   
   int h = ThrowHadronType(parton, scale);
   switch(h) {
      case  211:
      case -211: // Deliberate fall-through
      case  111: // Deliberate fall-through
         break;
      default:
         return mDss.Fragment(parton, scale);
         break;
   } // switch(h)
   
   Hadron hadron;
   hadron.type = h;
   
   SetFunctionParameters(h);
   
   // Find the proton beam spin in the parton coordinate frame.
   // Set function parameter 6 to the spin component transverse
   // to the parton direction.
   if(parton.cos_spin_depol > 0. and parton.cos_spin_depol <=1.) {
      Function()->SetParameter(6, parton.cos_spin_depol);
   }
   else {
      mFunction->SetParameter(6, 0.);
   }
   double z(-1.), phi(-1.);
   while(z > mZUpper or z < mZLower) {
      Function()->GetRandom2(z, phi);
   } // while
   hadron.z = z;
   hadron.phigen = phi;
   hadron.mParton = parton;
   
   static const double MEAN_JT2 = 0.2;
   double jt_sigma = sqrt(MEAN_JT2/2.0);  
   double jt; 
   // Throw jt according to an exponential distribution such
   // that mean jt^2 is 0.2
   do {
      jt = jt_sigma*sqrt( -2.0*log(gRandom->Uniform(1.0)) ); 	  
   } while(jt>(3.0*jt_sigma));
   
   // Randomly generate phi.
   TVector3 mom(jt, 0., 0.);
   mom.SetPhi(phi);
   // TVector3::SetPhi rotates phi leaving magnitude and theta unchanged.
   hadron.px = mom.Px();
   hadron.py = mom.Py();
   
   if(211 == h and UQUARK == parton.Pdg() and 1 == parton.Spin()) {
      cffPiU()->Fill(z, phi);
   } // if
   if(111 == h and DBAR == parton.Pdg() and -1 == parton.Spin()) {
      cffPi0DbarSd()->Fill(z, phi);
   } // if
   
   return std::vector<Hadron>(1, hadron);
}

TH2* CollinsFragmentation::cffPiU() {
   
   static TH2D* cff_pi_u(NULL);
   if(not cff_pi_u) {
      cff_pi_u = new TH2D("cff_pi_u",
                          "Collins FF as sampled (pi+ from u quark, spin up)",
                          100,0.0, 1.0,100,-(TMath::Pi()),TMath::Pi());
   } // if
   return cff_pi_u;
}

TH2* CollinsFragmentation::cffPi0DbarSd() {
   
   static TH2D* cff_pi0_dbar_sd(NULL);
   if(not cff_pi0_dbar_sd) {
      cff_pi0_dbar_sd = new TH2D("cff_pi0_dbar_sd",
                                 "Collins FF as sampled (pi0 from dbar quark, spin down)",
                                 100,0.0, 1.0,100,0.0,TMath::TwoPi());
   } // if
   return cff_pi0_dbar_sd;
}

TVector3
CollinsFragmentation::RotateToPartonFrame(const pythpart& parton,
                                          const TVector3& v) const {
   // Generate the rotation matrix that will put us into the 
   // fragmenting parton frame:
   
   TRotation rotM;
   rotM.SetZAxis(parton.PxPyPz());
   return rotM.Inverse() * v;
}

bool CollinsFragmentation::Supports(const pythpart& parton) const {
   return
   (abs(parton.Pdg()) < 4 and parton.Pdg() not_eq 0);
}

const std::vector<int>& CollinsFragmentation::Hadrons() const {
   static std::vector<int> hadrons;
   if(hadrons.empty()) {
      int array[3] = {
         211, -211, 111
      };
      hadrons.assign(array, array + 3);
   }
   return hadrons;
}

std::string CollinsFragmentation::Name() const {
   std::string name("DSS + Collins");
   if(mAtPositivityBound) {
      name.append( " (at positivity bound)");
   } // if
   return name;
}

} // namespace tppmc
